"""
* 名称：CobotGPT_V3副本_收集碰撞检测数据集.py
* 功能：收集碰撞检测模型的数据集：37列(关节角+骨骼坐标+碰撞标志)
* 其他：
"""

# from multiprocessing import Process
from multiprocessing import Process

import numpy as np
import torch
import torch.nn as nn

from Modules.KINECT.User.Releases.kinectcontrol import *
from Modules.AUBO.User.Releases.robotcontrol import *


def create_header():
    """
    * 功能：创建数据集.csv文件，仅运行一次即可
    * 其他：
    """
    with open('./logfiles/data_set_20240508.csv', 'a', newline = '', encoding = 'UTF-8') as f:
        关节角 = [f'Angle {i + 1}' for i in range(6)]
        骨骼坐标 = []
        是否碰撞 = ['collison']
        for i in range(10):
            骨骼坐标.extend([f'Joint {i} X'] + [f'Joint {i} Y'] + [f'Joint {i} Z'])
        header = 关节角 + 骨骼坐标 + 是否碰撞
        writer = csv.writer(f)
        writer.writerow(header)  # 写入表头
        print("创建数据集完毕，请立即注释此函数!/n")
        print("创建数据集完毕，请立即注释此函数!/n")
        print("创建数据集完毕，请立即注释此函数!/n")


class DNN(nn.Module):
    def __init__(self):
        """ 搭建神经网络各层 """
        super(DNN, self).__init__()
        self.net = nn.Sequential(  # 按顺序搭建各层
            nn.Linear(36, 32), nn.Sigmoid(),  # 第 1 层：全连接层
            nn.Linear(32, 8), nn.Sigmoid(),  # 第 2 层：全连接层
            nn.Linear(8, 4), nn.Sigmoid(),  # 第 3 层：全连接层
            nn.Linear(4, 1), nn.Sigmoid()  # 第 4 层：全连接层
        )

    def forward(self, x):
        """ 前向传播 """
        y = self.net(x)  # x 即输入数据
        return y  # y 即输出数据


class CobotGPT:

    def __init__(self):
        """
        功能：连接机械臂
        输入：
        输出：连接状态码
        """

        # 链接服务器
        # self.ip = '169.254.94.55'
        self.ip = '192.168.1.6'  # 注意配置网线口的适配器属性：关掉IPV6，手动设置IPV4地址和机械臂在同一个网段
        self.port = 8899
        self.joint_radian_list = [(80 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180),
                                  (
                                      160 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180,
                                      0 * pi / 180),
                                  (160 * pi / 180, -30 * pi / 180, 0 * pi / 180, 0 * pi / 180, 0 * pi / 180,
                                   0 * pi / 180),
                                  (160 * pi / 180, -30 * pi / 180, 80 * pi / 180, 0 * pi / 180, 0 * pi / 180,
                                   0 * pi / 180),
                                  (160 * pi / 180, -30 * pi / 180, 80 * pi / 180, 20 * pi / 180, 0 * pi / 180,
                                   0 * pi / 180),
                                  (160 * pi / 180, -30 * pi / 180, 80 * pi / 180, 20 * pi / 180, 90 * pi / 180,
                                   0 * pi / 180),  # 第一阶段
                                  (160 * pi / 180, 0 * pi / 180, 80 * pi / 180, 20 * pi / 180, 90 * pi / 180,
                                   0 * pi / 180),
                                  (90 * pi / 180, 0 * pi / 180, 80 * pi / 180, 20 * pi / 180, 90 * pi / 180,
                                   0 * pi / 180),
                                  (90 * pi / 180, -20 * pi / 180, 80 * pi / 180, 10 * pi / 180, 90 * pi / 180,
                                   0 * pi / 180)]
        # self.joint_radian_list = [
        #     (160 * pi / 180, 0 * pi / 180, 90 * pi / 180, 0 * pi / 180, 90 * pi / 180, -170 * pi / 180),
        #     (175 * pi / 180, 0 * pi / 180, 90 * pi / 180, 0 * pi / 180, 90 * pi / 180, -170 * pi / 180)]

    def robot_set_position(self, cnt, queue1: Queue, queue2: Queue):
        """
        功能：设置机械臂的固定轨迹点，并实时获取关节角信息
        输入：cnt输入测试次数，queue1输入是否碰撞，queue2输入机械臂关节角
        输出：
        """

        # 初始化logger
        logger_init()
        # 启动测试
        logger.info("{0} robot_set_position() running...".format(Auboi5Robot.get_local_time()))
        # 系统初始化
        Auboi5Robot.initialize()
        # 创建机械臂控制类
        robot = Auboi5Robot()
        # 创建上下文
        handle = robot.create_context()
        # 打印上下文
        logger.info("robot.rshd={0}".format(handle))

        try:
            # 连接服务器
            ip = self.ip
            port = self.port
            result = robot.connect(ip, port)
            if result != RobotErrorType.RobotError_SUCC:
                logger.info("connect server{0}:{1} failed.".format(ip, port))
            else:
                # 上电
                robot.robot_startup()
                # 设置碰撞等级
                robot.set_collision_class(10)
                # 初始化全局配置文件
                robot.init_profile()

                # 设置关节最大加速度
                robot.set_joint_maxacc((1.5, 1.5, 1.5, 1.5, 1.5, 1.5))
                # 设置关节最大速度
                # robot.set_joint_maxvelc((0.2, 0.2, 0.2, 0.2, 0.2, 0.2))
                robot.set_joint_maxvelc((0.1, 0.1, 0.1, 0.1, 0.1, 0.1))  # 低速
                # robot.set_joint_maxvelc((0.3, 0.3, 0.3, 0.3, 0.3, 0.3))

                # 回到初始位姿
                joint_radian_ori = (0, 0, 0, 0, 0, 0)  # 初始位姿,tuple类型
                logger.info("回到初始位姿 {0}".format(joint_radian_ori))
                robot.move_joint(joint_radian_ori)

                # 循环测试
                while cnt > 0:
                    cnt -= 1

                    for i in range(len(self.joint_radian_list)):
                        joint_radian = self.joint_radian_list[i]
                        logger.info("move joint to {0}".format(joint_radian))
                        robot.move_joint(joint_radian)
                        # print(f"i:{i}")


        # 出现错误的处理
        except RobotError as e:
            logger.error("{0} robot Event:{1} in robot_set_position()".format(robot.get_local_time(), e))

        except KeyboardInterrupt:
            # 断开服务器链接
            if robot.connected:
                # 断开机械臂链接
                robot.disconnect()
            # 释放库资源
            Auboi5Robot.uninitialize()
            print("robot disconnected-------------------------")

        # 无论是否有异常，都会执行的代码
        finally:
            # 断开服务器链接
            # if robot.connected:
            #     # 关闭机械臂
            #     robot.robot_shutdown()
            #     # 断开机械臂链接
            #     robot.disconnect()
            # # 释放库资源
            # Auboi5Robot.uninitialize()
            logger.info("{0} test completed.".format(Auboi5Robot.get_local_time()))

    def dataset(self, queue1: Queue, queue2: Queue):
        """
                功能：将坐标数据转换成数据集
                输入：queue1:输入骨骼坐标，queue2：输出36列坐标
                输出：
                """

        # 初始化logger
        logger_init()
        # 启动测试
        logger.info("{0} dataset() running...".format(Auboi5Robot.get_local_time()))
        # 系统初始化
        Auboi5Robot.initialize()
        # 创建机械臂控制类
        robot = Auboi5Robot()
        # 创建上下文
        handle = robot.create_context()
        # 打印上下文
        logger.info("robot.rshd={0}".format(handle))

        try:
            # 连接服务器
            ip = self.ip
            port = self.port
            result = robot.connect(ip, port)
            if result != RobotErrorType.RobotError_SUCC:
                logger.info("connect server{0}:{1} failed.".format(ip, port))
            else:
                # 先创建数据集，其它代码给我注释喽
                # create_header()

                # 创建完数据集后，再运行以下代码
                with open('./logfiles/data_set_20240508.csv', 'a', newline = '', encoding = 'UTF-8') as f:
                    # 关节角 = [f'Angle {i + 1}' for i in range(6)]
                    # 骨骼坐标 = []
                    # 是否碰撞 = ['collison']
                    # for i in range(10):
                    #     骨骼坐标.extend([f'Joint {i} X'] + [f'Joint {i} Y'] + [f'Joint {i} Z'])
                    # header = 关节角 + 骨骼坐标 + 是否碰撞
                    writer = csv.writer(f)
                    # writer.writerow(header)  # 写入表头

                    while result == RobotErrorType.RobotError_SUCC:
                        current_waypoint = robot.get_current_waypoint()
                        data1 = current_waypoint["joint"]  # 机器人当前关节角：6
                        data2 = queue1.get()    # 当前人体骨骼坐标：30
                        collison_flag = [1]     # 0：当前采集的是非碰撞的数据集
                        # print(f"data1:{data1}") # 机械臂关节角，6列
                        # print(f"data2:{data2}")  # 人体骨骼坐标，3*10=30列
                        data = data1 + data2 + collison_flag    # 37元素的列表
                        if int(len(data)) == 37:
                            print(data)
                            writer.writerow(data)

        # 出现错误的处理
        except RobotError as e:
            logger.error("{0} robot Event:{1}".format(self.robot.get_local_time(), e))

    def collision_detection(self, queue1: Queue, queue2: Queue, queue3: Queue):
        """
        功能：碰撞检测
        输入：queue1:输出是否碰撞，queue2:输出关节角，queue3：输入36列坐标
        输出：
        """

        collision_flag = 0  # 默认未发生碰撞
        step = 0
        model = torch.load('DNN_model.pth')  # 加载训练好的模型
        model.to('cuda:0')

        while True:

            while queue1.empty() is not True:
                queue1.get()
            while queue2.empty() is not True:
                queue2.get()
            if collision_flag == 0:
                # print("collision 运行正常传输关节角")
                queue1.put(0)
                queue2.put(self.joint_radian_list[step])
                step += 1
                if step == int(len(self.joint_radian_list)):
                    step = 0
            else:
                print("collision 传输碰撞信号")
                # collision_flag = 0
                queue1.put(1)

            X = queue3.get()  # queue3：输入总坐标，列表
            # print(f"self.data:{self.data}")
            # print(len(X))
            # print(X)
            if len(X) == 36:
                arr = np.array(X)  # 列表 => numpy数组
                arr = arr.astype(np.float32)  # 转为 float32 类型数组
                ts = torch.tensor(arr)  # numpy数组 => 张量
                ts = ts.to('cuda')
                with torch.no_grad():  # 该局部关闭梯度计算功能
                    pred = model(ts)  # 一次前向传播（批量）
                    pred[pred >= 0.5] = 1
                    pred[pred < 0.5] = 0
                    pred = int(pred.item())  # 张量 => int值
                    print(f"是否碰撞: {pred}")
            else:
                # print("输入的总坐标格式错误")
                pred = 0
            if pred == 1:
                collision_flag = 1
            else:
                collision_flag = 0
            # time.sleep(0.5)


if __name__ == '__main__':
    # 创建类实例
    game = BodyGameRuntime()
    cobot = CobotGPT()
    # 创建通信队列
    q1 = Queue(maxsize = 1)  # q1:是否碰撞，
    q2 = Queue(maxsize = 1)  # q2:机械臂关节角
    q3 = Queue()  # q3:kinect_waypoint
    q4 = Queue(maxsize = 1)  # q4:36列坐标

    # # 创建子进程
    p1 = Process(target = cobot.robot_set_position, kwargs = {"cnt": 5, "queue1": q1, "queue2": q2})
    # p2 = Process(target=cobot.collision_detection, kwargs={"queue1": q1, "queue2": q2, "queue3": q4})
    p3 = Process(target = cobot.dataset, kwargs = {"queue1": q3, "queue2": q4})
    p1.start()
    # p2.start()
    p3.start()
    # 主进程
    game.run(q3)
